import rclpy
from rclpy.node import Node
from sensor_msgs.msg import Imu
from sls_msgs.msg import LlaVelocity
from geometry_msgs.msg import Quaternion
from scipy.spatial.transform import Rotation
import numpy as np
import math

class GPSKalmanFilter:
    def __init__(self, initial_state, initial_covariance, process_noise, measurement_noise):
        self.x = initial_state
        self.P = initial_covariance
        self.Q = process_noise
        self.R = measurement_noise
        self.F = np.eye(6)
        self.H = np.array([
            [1, 0, 0, 0, 0, 0],
            [0, 1, 0, 0, 0, 0],
            [0, 0, 1, 0, 0, 0],
            [0, 0, 0, 1, 0, 0]
        ])

    def predict(self, dt):
        vel = self.x[3]
        yaw = self.x[5]
        
        # 1m
        lat_scale = 1.0 / 111320.0
        lon_scale = 1.0 / (111320.0 * math.cos(math.radians(self.x[0])))
        
        self.F[0, 3] = dt * math.cos(yaw) * lat_scale
        self.F[1, 3] = dt * math.sin(yaw) * lon_scale
        self.F[4, 5] = 1.0
        
        # 预测步骤
        self.x = self.F @ self.x
        self.P = self.F @ self.P @ self.F.T + self.Q

    def update(self, z):
        # 计算卡尔曼增益
        y = z - self.H @ self.x
        S = self.H @ self.P @ self.H.T + self.R
        K = self.P @ self.H.T @ np.linalg.inv(S)
        
        # 更新状态和协方差
        self.x = self.x + K @ y
        self.P = (np.eye(6) - K @ self.H) @ self.P

class LlaKalmanInterpolator(Node):
    def __init__(self):
        super().__init__('lla_kalman_interpolator')
        
        # 订阅IMU数据 (高频)
        self.imu_sub = self.create_subscription(
            Imu,
            'imu/data',
            self.imu_callback,
            10)
        
        # 订阅LLA速度信息 (低频)
        self.lla_vel_sub = self.create_subscription(
            LlaVelocity,
            'lla_velocity',
            self.lla_vel_callback,
            10)
        
        # 初始化卡尔曼滤波器
        initial_state = np.zeros(6)
        initial_covariance = np.diag([1.0, 1.0, 10.0, 1.0, 0.1, 0.1])
        process_noise = np.diag([1e-6, 1e-6, 1e-4, 1e-4, 1e-4, 1e-4])
        measurement_noise = np.diag([1e-8, 1e-8, 1e-2, 1e-4])
        
        self.kf = GPSKalmanFilter(
            initial_state,
            initial_covariance,
            process_noise,
            measurement_noise
        )
        
        # 状态变量
        self.last_imu_time = None
        self.initialized = False
        
        # 插值结果发布
        self.interpolated_pub = self.create_publisher(LlaVelocity, 'lla_interpolated', 10)
        
        self.get_logger().info("LLA Kalman Interpolator Node Initialized")

    def initialize_filter(self, msg):
        """使用第一个GPS消息初始化滤波器"""
        self.kf.x = np.array([
            msg.latitude, 
            msg.longitude, 
            msg.altitude,
            msg.vel,
            msg.vel_heading,
            msg.yaw if not math.isnan(msg.yaw) else 0.0
        ])
        self.initialized = True
        self.get_logger().info("Kalman Filter Initialized")

    def imu_callback(self, msg):
        if not self.initialized:
            return
            
        # 从IMU数据提取yaw角
        quat = msg.orientation
        rotation = Rotation.from_quat([quat.x, quat.y, quat.z, quat.w])
        yaw = rotation.as_euler('zyx')[0]  # 弧度
        
        # 更新时间差
        now = self.get_clock().now()
        if self.last_imu_time is None:
            self.last_imu_time = now
            self.kf.x[5] = yaw  # 初始化yaw
            return
            
        dt = (now - self.last_imu_time).nanoseconds * 1e-9
        self.last_imu_time = now
        
        # 更新yaw状态
        self.kf.x[5] = yaw
        
        # 执行预测步骤
        self.kf.predict(dt)
        
        # 发布插值结果
        self.publish_interpolated()
    
    def lla_vel_callback(self, msg):
        if not self.initialized:
            self.initialize_filter(msg)
            return
            
        # 准备测量向量
        z = np.array([msg.latitude, msg.longitude, msg.altitude, msg.vel])
        
        # 执行更新步骤
        self.kf.update(z)
        
        # 如果GPS提供了yaw，可以选择性地更新
        if not math.isnan(msg.yaw):
            self.kf.x[5] = msg.yaw
        
        # 发布更新后的状态
        self.publish_interpolated()
    
    def publish_interpolated(self):
        """发布当前卡尔曼滤波状态"""
        interpolated_msg = LlaVelocity()
        interpolated_msg.latitude = self.kf.x[0]
        interpolated_msg.longitude = self.kf.x[1]
        interpolated_msg.altitude = self.kf.x[2]
        interpolated_msg.vel = self.kf.x[3]
        interpolated_msg.vel_heading = self.kf.x[4]
        interpolated_msg.yaw = self.kf.x[5]
        
        self.interpolated_pub.publish(interpolated_msg)

def main(args=None):
    rclpy.init(args=args)
    node = LlaKalmanInterpolator()
    rclpy.spin(node)
    node.destroy_node()
    rclpy.shutdown()

if __name__ == '__main__':
    main()